import numpy as np
import cv2
import json
import os
import logging
import pandas as pd
import easydict as edict

from pose_handler import Test
from pose_handler import pos2vec

def test_velocity_visualize():
    import matplotlib.pyplot as plt
    import matplotlib
    # matplotlib.use('Qt5Agg')
    # seq_name = 'MOT16-02'
    # json_path = 'alphapose_tf_results/result_mot16_02/POSE/alpha-pose-results.json'
    # my_mot_det_file = './output/mymot_joints_det_{}.txt'.format(seq_name)

    # seq_nqme = 'TUD-Stadtmitte'
    seq_nqme = 'TUD-Campus】'
    json_path = './output/mot15_pose/train/{}/POSE/alpha-pose-results.json'.format(seq_nqme)
    test = Test(json_path)

    if True:
        # test.load_gt('./mot16Label/train/MOT16-02/gt/gt.txt')
        test.load_gt('MOT15/train/{}/gt/gt.txt'.format(seq_nqme))
        test.parse_xml()
        test.attach_skeleton()
        df = test.get_skeleton_data()
        df.to_hdf('./output/foo.h5', 'df')
    else:
        test.load_df('./output/foo.h5')
    tid = 5
    color = ['r', 'g', 'b']
    # m_var = pd.DataFrame(None, columns=['tid', 'joint_id', 'var'])
    m_var = pd.DataFrame(np.zeros((1, 15)), columns=['j' + str(i) for i in range(15)])
    # output_file = open('./output/joint_stability.txt', 'w')

    df = test.df[test.df.tid == tid]

    if len(df) == 0:
        logging.error('no this target:' + str(tid))
    assert len(df) > 0
    df = df.drop_duplicates('fid', 'first', False)

    for joint_index in range(15):
        # pos, fid = test.get_vec_material(joint_index, tid)

        pos = df.loc[:, ['x' + str(joint_index), 'y' + str(joint_index)]].values
        fid = df.loc[:, 'fid'].values
        assert len(fid) == len(pos)
        print(fid)
        vec = pos2vec(fid, pos)
        # todo 归一化。
        plt.plot(np.ones(len(vec)) * joint_index, vec, color[joint_index % 3] + 'o')
        var = np.var(vec)
        m_var.loc[0, 'j' + str(joint_index)] = var
        # print("%d, %d, %d" %( tid, joint_index, var), file=output_file)
    m_var.to_csv('./output/joint_stability_{}.txt'.format(seq_nqme))
    values = m_var.as_matrix()
    print(values)
    print('tid %d , joint : rank' % tid)
    print(np.argsort(values[0]))
    # plt.axis([-1, 15, 0, 30])
    # plt.show()


if __name__ == '__main__':
    test_velocity_visualize()